package com.example.flink;

import org.apache.flink.streaming.api.functions.sink.RichSinkFunction;

import com.esotericsoftware.minlog.Log;

import lombok.extern.slf4j.Slf4j;

import org.apache.flink.configuration.Configuration;
import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.PreparedStatement;
import java.util.ArrayList;
import java.util.List;

@Slf4j
public class ClickHouseSink extends RichSinkFunction<LogEvent> {
    private transient Connection connection;
    private transient PreparedStatement ps;
    private transient List<LogEvent> buffer;
    private final int batchSize = 100;
    private final long flushIntervalMs = 5000;
    private transient long lastFlushTime;

    @Override
    public void open(Configuration parameters) throws Exception {
        log.info("ClickHouseSink open");
        connection = DriverManager.getConnection("jdbc:clickhouse://localhost:8123/default", "default", "");
        ps = connection.prepareStatement(
            "INSERT INTO cleaned_app_log (timestamp, level, message, thread, logger, trace_id, span_id) VALUES (?, ?, ?, ?, ?, ?, ?)"
        );
        buffer = new ArrayList<>();
        lastFlushTime = System.currentTimeMillis();
    }

    @Override
    public void invoke(LogEvent event, Context context) {
        try {
            buffer.add(event);
            long now = System.currentTimeMillis();
            if (buffer.size() >= batchSize || (now - lastFlushTime) >= flushIntervalMs) {
                for (LogEvent e : buffer) {
                    String tsStr = e.timestamp;
                    String formattedTime = null;
                    if (tsStr != null && !tsStr.isEmpty()) {
                        long ts = Long.parseLong(tsStr);
                        formattedTime = java.time.LocalDateTime.ofInstant(
                            java.time.Instant.ofEpochMilli(ts),
                            java.time.ZoneId.systemDefault()
                        ).toString().replace('T', ' ').substring(0, 19);
                    }
                    ps.setString(1, formattedTime);
                    ps.setString(2, e.level);
                    ps.setString(3, e.message);
                    ps.setString(4, e.thread);
                    ps.setString(5, e.logger);
                    ps.setString(6, e.trace_id);
                    ps.setString(7, e.span_id);
                    ps.addBatch();
                }
                ps.executeBatch();
                buffer.clear();
                lastFlushTime = now;
            }
        } catch (Exception e) {
            log.error("ClickHouseSink error", e);
            e.printStackTrace();
        }
    }

    @Override
    public void close() throws Exception {
        if (buffer != null && !buffer.isEmpty()) {
            for (LogEvent e : buffer) {
                String tsStr = e.timestamp;
                String formattedTime = null;
                if (tsStr != null && !tsStr.isEmpty()) {
                    long ts = Long.parseLong(tsStr);
                    formattedTime = java.time.LocalDateTime.ofInstant(
                        java.time.Instant.ofEpochMilli(ts),
                        java.time.ZoneId.systemDefault()
                    ).toString().replace('T', ' ').substring(0, 19);
                }
                ps.setString(1, formattedTime);
                ps.setString(2, e.level);
                ps.setString(3, e.message);
                ps.setString(4, e.thread);
                ps.setString(5, e.logger);
                ps.setString(6, e.trace_id);
                ps.setString(7, e.span_id);
                ps.addBatch();
            }
            ps.executeBatch();
            buffer.clear();
        }
        if (ps != null) ps.close();
        if (connection != null) connection.close();
    }
} 